/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.io.gcp.bigquery;



import com.google.api.client.util.BackOff;

import com.google.api.client.util.BackOffUtils;

import com.google.api.client.util.Sleeper;

import com.google.api.services.bigquery.model.*;

import com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto;

import com.google.cloud.hadoop.util.ApiErrorExtractor;

import com.bff.gaia.unified.sdk.extensions.gcp.util.BackOffAdapter;

import com.bff.gaia.unified.sdk.io.FileSystems;

import com.bff.gaia.unified.sdk.io.fs.ResolveOptions;

import com.bff.gaia.unified.sdk.options.ValueProvider;

import com.bff.gaia.unified.sdk.options.ValueProvider.NestedValueProvider;

import com.bff.gaia.unified.sdk.transforms.SerializableFunction;

import com.bff.gaia.unified.sdk.util.FluentBackoff;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import com.bff.gaia.unified.vendor.guava.com.google.common.hash.Hashing;

import org.joda.time.Duration;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import javax.annotation.Nullable;

import java.io.IOException;

import java.util.ArrayList;

import java.util.List;

import java.util.Map;

import java.util.UUID;

import java.util.regex.Matcher;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkState;



/** A set of helper functions and classes used by {@link BigQueryIO}. */

public class BigQueryHelpers {

  private static final String RESOURCE_NOT_FOUND_ERROR =

      "BigQuery %1$s not found for table \"%2$s\" . Please create the %1$s before pipeline"

          + " execution. If the %1$s is created by an earlier stage of the pipeline, this"

          + " validation can be disabled using #withoutValidation.";



  private static final String UNABLE_TO_CONFIRM_PRESENCE_OF_RESOURCE_ERROR =

      "Unable to confirm BigQuery %1$s presence for table \"%2$s\". If the %1$s is created by"

          + " an earlier stage of the pipeline, this validation can be disabled using"

          + " #withoutValidation.";



  private static final Logger LOG = LoggerFactory.getLogger(BigQueryHelpers.class);



  // Given a potential failure and a current job-id, return the next job-id to be used on retry.

  // Algorithm is as follows (given input of job_id_prefix-N)

  //   If BigQuery has no status for job_id_prefix-n, we should retry with the same id.

  //   If job-id-prefix-n is in the PENDING or successful states, no retry is needed.

  //   Otherwise (job-id-prefix-n completed with errors), try again with job-id-prefix-(n+1)

  //

  // We continue to loop through these job ids until we find one that has either succeed, or that

  // has not been issued yet.

  static class RetryJobIdResult {

    public final RetryJobId jobId;

    public final boolean shouldRetry;



    public RetryJobIdResult(RetryJobId jobId, boolean shouldRetry) {

      this.jobId = jobId;

      this.shouldRetry = shouldRetry;

    }

  }



  // A class that waits for pending jobs, retrying them according to policy if they fail.

  static class PendingJobManager {

    private static class JobInfo {

      private final PendingJob pendingJob;

      @Nullable private final SerializableFunction<PendingJob, Exception> onSuccess;



      public JobInfo(PendingJob pendingJob, SerializableFunction<PendingJob, Exception> onSuccess) {

        this.pendingJob = pendingJob;

        this.onSuccess = onSuccess;

      }

    }



    private List<JobInfo> pendingJobs = Lists.newArrayList();

    private final BackOff backOff;



    PendingJobManager() {

      this(

          BackOffAdapter.toGcpBackOff(

              FluentBackoff.DEFAULT

                  .withMaxRetries(Integer.MAX_VALUE)

                  .withInitialBackoff(Duration.standardSeconds(1))

                  .withMaxBackoff(Duration.standardMinutes(1))

                  .backoff()));

    }



    PendingJobManager(BackOff backOff) {

      this.backOff = backOff;

    }



    // Add a pending job and a function to call when the job has completed successfully.

    PendingJobManager addPendingJob(

        PendingJob pendingJob, @Nullable SerializableFunction<PendingJob, Exception> onSuccess) {

      this.pendingJobs.add(new JobInfo(pendingJob, onSuccess));

      return this;

    }



    void waitForDone() throws Exception {

      LOG.info("Waiting for jobs to complete.");

      Sleeper sleeper = Sleeper.DEFAULT;

      while (!pendingJobs.isEmpty()) {

        List<JobInfo> retryJobs = Lists.newArrayList();

        for (JobInfo jobInfo : pendingJobs) {

          if (jobInfo.pendingJob.pollJob()) {

            // Job has completed successfully.

            LOG.info("Job {} completed successfully.", jobInfo.pendingJob.currentJobId);

            Exception e = jobInfo.onSuccess.apply(jobInfo.pendingJob);

            if (e != null) {

              throw e;

            }

          } else {

            // Job failed, schedule it again.

            LOG.info("Job {} failed. retrying.", jobInfo.pendingJob.currentJobId);

            retryJobs.add(jobInfo);

          }

        }

        pendingJobs = retryJobs;

        if (!pendingJobs.isEmpty()) {

          // Sleep before retrying.

          nextBackOff(sleeper, backOff);

          // Run the jobs to retry. If a job has hit the maximum number of retries then runJob

          // will raise an exception.

          for (JobInfo job : pendingJobs) {

            job.pendingJob.runJob();

          }

        }

      }

    }



    /** Identical to {@link BackOffUtils#next} but without checked IOException. */

    private static boolean nextBackOff(Sleeper sleeper, BackOff backOff)

        throws InterruptedException {

      try {

        return BackOffUtils.next(sleeper, backOff);

      } catch (IOException e) {

        throw new RuntimeException(e);

      }

    }

  }



  static class PendingJob {

    private final SerializableFunction<RetryJobId, Void> executeJob;

    private final SerializableFunction<RetryJobId, Job> pollJob;

    private final SerializableFunction<RetryJobId, Job> lookupJob;

    private final int maxRetries;

    private int currentAttempt;

    RetryJobId currentJobId;

    Job lastJobAttempted;

    boolean started;



    PendingJob(

        SerializableFunction<RetryJobId, Void> executeJob,

        SerializableFunction<RetryJobId, Job> pollJob,

        SerializableFunction<RetryJobId, Job> lookupJob,

        int maxRetries,

        String jobIdPrefix) {

      this.executeJob = executeJob;

      this.pollJob = pollJob;

      this.lookupJob = lookupJob;

      this.maxRetries = maxRetries;

      this.currentAttempt = 0;

      currentJobId = new RetryJobId(jobIdPrefix, 0);

      this.started = false;

    }



    // Run the job.

    void runJob() throws IOException {

      ++currentAttempt;

      if (!shouldRetry()) {

        throw new RuntimeException(

            String.format(

                "Failed to create job with prefix %s, "

                    + "reached max retries: %d, last failed job: %s.",

                currentJobId.getJobIdPrefix(),

                maxRetries,

                BigQueryHelpers.jobToPrettyString(lastJobAttempted)));

      }



      try {

        this.started = false;

        executeJob.apply(currentJobId);

      } catch (RuntimeException e) {

        LOG.warn("Job {} failed with {}", currentJobId.getJobId(), e);

        // It's possible that the job actually made it to BQ even though we got a failure here.

        // For example, the response from BQ may have timed out returning. getRetryJobId will

        // return the correct job id to use on retry, or a job id to continue polling (if it turns

        // out that the job has not actually failed yet).

        RetryJobIdResult result = getRetryJobId(currentJobId, lookupJob);

        currentJobId = result.jobId;

        if (result.shouldRetry) {

          // Otherwise the jobs either never started or started and failed. Try the job again with

          // the job id returned by getRetryJobId.

          LOG.info("Will retry with job id {}", currentJobId.getJobId());

          return;

        }

      }

      LOG.info("job {} started", currentJobId.getJobId());

      // The job has reached BigQuery and is in either the PENDING state or has completed

      // successfully.

      this.started = true;

    }



    // Poll the status of the job. Returns true if the job has completed successfully and false

    // otherwise.

    boolean pollJob() throws IOException {

      if (started) {

        Job job = pollJob.apply(currentJobId);

        this.lastJobAttempted = job;

        Status jobStatus = parseStatus(job);

        switch (jobStatus) {

          case SUCCEEDED:

            LOG.info("Load job {} succeeded. Statistics: {}", currentJobId, job.getStatistics());

            return true;

          case UNKNOWN:

            // This might happen if BigQuery's job listing is slow. Retry with the same

            // job id.

            LOG.info(

                "Load job {} finished in unknown state: {}: {}",

                currentJobId,

                job.getStatus(),

                shouldRetry() ? "will retry" : "will not retry");

            return false;

          case FAILED:

            String oldJobId = currentJobId.getJobId();

            currentJobId = BigQueryHelpers.getRetryJobId(currentJobId, lookupJob).jobId;

            LOG.info(

                "Load job {} failed, {}: {}. Next job id {}",

                oldJobId,

                shouldRetry() ? "will retry" : "will not retry",

                job.getStatus(),

                currentJobId);

            return false;

          default:

            throw new IllegalStateException(

                String.format(

                    "Unexpected status [%s] of load job: %s.",

                    job.getStatus(), BigQueryHelpers.jobToPrettyString(job)));

        }

      }

      return false;

    }



    boolean shouldRetry() {

      return currentAttempt < maxRetries + 1;

    }

  }



  static class RetryJobId {

    private final String jobIdPrefix;

    private final int retryIndex;



    RetryJobId(String jobIdPrefix, int retryIndex) {

      this.jobIdPrefix = jobIdPrefix;

      this.retryIndex = retryIndex;

    }



    String getJobIdPrefix() {

      return jobIdPrefix;

    }



    int getRetryIndex() {

      return retryIndex;

    }



    String getJobId() {

      return jobIdPrefix + "-" + retryIndex;

    }



    @Override

    public String toString() {

      return getJobId();

    }

  }



  static RetryJobIdResult getRetryJobId(

      RetryJobId currentJobId, SerializableFunction<RetryJobId, Job> lookupJob) {

    for (int retryIndex = currentJobId.getRetryIndex(); ; retryIndex++) {

      RetryJobId jobId = new RetryJobId(currentJobId.getJobIdPrefix(), retryIndex);

      try {

        Job loadJob = lookupJob.apply(jobId);

        if (loadJob == null) {

          LOG.info("job id {} not found, so retrying with that id", jobId);

          // This either means that the original job was never properly issued (on the first

          // iteration of the loop) or that we've found a retry id that has not been used yet. Try

          // again with this job id.

          return new RetryJobIdResult(jobId, true);

        }

        JobStatus jobStatus = loadJob.getStatus();

        if (jobStatus == null) {

          LOG.info("job status for {} not found, so retrying with that job id", jobId);

          return new RetryJobIdResult(jobId, true);

        }

        if ("PENDING".equals(jobStatus.getState()) || "RUNNING".equals(jobStatus.getState())) {

          // The job id has been issued and is currently pending. This can happen after receiving

          // an error from the load or copy job creation (e.g. that error might come because the

          // job already exists). Return to the caller which job id is pending (it might not be the

          // one passed in) so the caller can then wait for this job to finish.

          LOG.info("job {} in pending or running state, so continuing with that job id", jobId);

          return new RetryJobIdResult(jobId, false);

        }

        if (jobStatus.getErrorResult() == null

            && (jobStatus.getErrors() == null || jobStatus.getErrors().isEmpty())) {

          // Import succeeded. No retry needed.

          LOG.info("job {} succeeded, so not retrying ", jobId);

          return new RetryJobIdResult(jobId, false);

        }

        // This job has failed, so we assume the data cannot enter BigQuery. We will check the next

        // job in the sequence (with the same unique prefix) to see if is either pending/succeeded

        // or can be used to generate a retry job.

        LOG.info("job {} is failed. Checking the next job id", jobId);

      } catch (RuntimeException e) {

        LOG.info("caught exception while querying job {}", jobId);

        return new RetryJobIdResult(jobId, true);

      }

    }

  }



  /** Status of a BigQuery job or request. */

  enum Status {

    SUCCEEDED,

    FAILED,

    UNKNOWN,

  }



  @VisibleForTesting

  static TableReferenceProto.TableReference toTableRefProto(TableReference ref) {

    TableReferenceProto.TableReference.Builder builder =

        TableReferenceProto.TableReference.newBuilder();

    if (ref.getProjectId() != null) {

      builder.setProjectId(ref.getProjectId());

    }

    return builder.setDatasetId(ref.getDatasetId()).setTableId(ref.getTableId()).build();

  }



  @VisibleForTesting

  static TableReference toTableRef(TableReferenceProto.TableReference ref) {

    return new TableReference()

        .setProjectId(ref.getProjectId())

        .setDatasetId(ref.getDatasetId())

        .setTableId(ref.getTableId());

  }



  /** Return a displayable string representation for a {@link TableReference}. */

  @Nullable

  static ValueProvider<String> displayTable(@Nullable ValueProvider<TableReference> table) {

    if (table == null) {

      return null;

    }

    return NestedValueProvider.of(table, new TableRefToTableSpec());

  }



  /** Returns a canonical string representation of the {@link TableReference}. */

  public static String toTableSpec(TableReference ref) {

    StringBuilder sb = new StringBuilder();

    if (ref.getProjectId() != null) {

      sb.append(ref.getProjectId());

      sb.append(":");

    }



    sb.append(ref.getDatasetId()).append('.').append(ref.getTableId());

    return sb.toString();

  }



  @Nullable

  static ValueProvider<String> displayTableRefProto(

      @Nullable ValueProvider<TableReferenceProto.TableReference> table) {

    if (table == null) {

      return null;

    }



    return NestedValueProvider.of(table, new TableRefProtoToTableSpec());

  }



  /** Returns a canonical string representation of a {@link TableReferenceProto.TableReference}. */

  public static String toTableSpec(TableReferenceProto.TableReference ref) {

    StringBuilder sb = new StringBuilder();

    if (ref.getProjectId() != null) {

      sb.append(ref.getProjectId());

      sb.append(":");

    }



    sb.append(ref.getDatasetId()).append('.').append(ref.getTableId());

    return sb.toString();

  }



  static <K, V> List<V> getOrCreateMapListValue(Map<K, List<V>> map, K key) {

    return map.computeIfAbsent(key, k -> new ArrayList<>());

  }



  /**

   * Parse a table specification in the form {@code "[project_id]:[dataset_id].[table_id]"} or

   * {@code "[dataset_id].[table_id]"}.

   *

   * <p>If the project id is omitted, the default project id is used.

   */

  public static TableReference parseTableSpec(String tableSpec) {

    Matcher match = BigQueryIO.TABLE_SPEC.matcher(tableSpec);

    if (!match.matches()) {

      throw new IllegalArgumentException(

          "Table reference is not in [project_id]:[dataset_id].[table_id] "

              + "format: "

              + tableSpec);

    }



    TableReference ref = new TableReference();

    ref.setProjectId(match.group("PROJECT"));



    return ref.setDatasetId(match.group("DATASET")).setTableId(match.group("TABLE"));

  }



  /** Strip off any partition decorator information from a tablespec. */

  public static String stripPartitionDecorator(String tableSpec) {

    int index = tableSpec.lastIndexOf('$');

    return (index == -1) ? tableSpec : tableSpec.substring(0, index);

  }



  static String jobToPrettyString(@Nullable Job job) throws IOException {

    if (job != null && job.getConfiguration().getLoad() != null) {

      // Removing schema and sourceUris from error messages for load jobs since these fields can be

      // quite long and error message might not be displayed properly in runner specific logs.

      job = job.clone();

      job.getConfiguration().getLoad().setSchema(null);

      job.getConfiguration().getLoad().setSourceUris(null);

    }



    return job == null ? "null" : job.toPrettyString();

  }



  static String statusToPrettyString(@Nullable JobStatus status) throws IOException {

    return status == null ? "Unknown status: null." : status.toPrettyString();

  }



  static Status parseStatus(@Nullable Job job) {

    if (job == null) {

      return Status.UNKNOWN;

    }

    JobStatus status = job.getStatus();

    if (status.getErrorResult() != null) {

      return Status.FAILED;

    } else if (status.getErrors() != null && !status.getErrors().isEmpty()) {

      return Status.FAILED;

    } else {

      return Status.SUCCEEDED;

    }

  }



  @VisibleForTesting

  static String toJsonString(Object item) {

    if (item == null) {

      return null;

    }

    try {

      return BigQueryIO.JSON_FACTORY.toString(item);

    } catch (IOException e) {

      throw new RuntimeException(

          String.format("Cannot serialize %s to a JSON string.", item.getClass().getSimpleName()),

          e);

    }

  }



  @VisibleForTesting

  static <T> T fromJsonString(String json, Class<T> clazz) {

    if (json == null) {

      return null;

    }

    try {

      return BigQueryIO.JSON_FACTORY.fromString(json, clazz);

    } catch (IOException e) {

      throw new RuntimeException(

          String.format("Cannot deserialize %s from a JSON string: %s.", clazz, json), e);

    }

  }



  /**

   * Returns a randomUUID string.

   *

   * <p>{@code '-'} is removed because BigQuery doesn't allow it in dataset id.

   */

  static String randomUUIDString() {

    return UUID.randomUUID().toString().replaceAll("-", "");

  }



  static void verifyTableNotExistOrEmpty(BigQueryServices.DatasetService datasetService, TableReference tableRef) {

    try {

      if (datasetService.getTable(tableRef) != null) {

        checkState(

            datasetService.isTableEmpty(tableRef),

            "BigQuery table is not empty: %s.",

            toTableSpec(tableRef));

      }

    } catch (IOException | InterruptedException e) {

      if (e instanceof InterruptedException) {

        Thread.currentThread().interrupt();

      }

      throw new RuntimeException(

          "unable to confirm BigQuery table emptiness for table " + toTableSpec(tableRef), e);

    }

  }



  static void verifyDatasetPresence(BigQueryServices.DatasetService datasetService, TableReference table) {

    try {

      datasetService.getDataset(table.getProjectId(), table.getDatasetId());

    } catch (Exception e) {

      ApiErrorExtractor errorExtractor = new ApiErrorExtractor();

      if ((e instanceof IOException) && errorExtractor.itemNotFound((IOException) e)) {

        throw new IllegalArgumentException(

            String.format(RESOURCE_NOT_FOUND_ERROR, "dataset", toTableSpec(table)), e);

      } else if (e instanceof RuntimeException) {

        throw (RuntimeException) e;

      } else {

        throw new RuntimeException(

            String.format(

                UNABLE_TO_CONFIRM_PRESENCE_OF_RESOURCE_ERROR, "dataset", toTableSpec(table)),

            e);

      }

    }

  }



  static String getDatasetLocation(

	  BigQueryServices.DatasetService datasetService, String projectId, String datasetId) {

    Dataset dataset;

    try {

      dataset = datasetService.getDataset(projectId, datasetId);

    } catch (Exception e) {

      if (e instanceof InterruptedException) {

        Thread.currentThread().interrupt();

      }

      throw new RuntimeException(

          String.format(

              "unable to obtain dataset for dataset %s in project %s", datasetId, projectId),

          e);

    }

    return dataset.getLocation();

  }



  static void verifyTablePresence(BigQueryServices.DatasetService datasetService, TableReference table) {

    try {

      datasetService.getTable(table);

    } catch (Exception e) {

      ApiErrorExtractor errorExtractor = new ApiErrorExtractor();

      if ((e instanceof IOException) && errorExtractor.itemNotFound((IOException) e)) {

        throw new IllegalArgumentException(

            String.format(RESOURCE_NOT_FOUND_ERROR, "table", toTableSpec(table)), e);

      } else if (e instanceof RuntimeException) {

        throw (RuntimeException) e;

      } else {

        throw new RuntimeException(

            String.format(

                UNABLE_TO_CONFIRM_PRESENCE_OF_RESOURCE_ERROR, "table", toTableSpec(table)),

            e);

      }

    }

  }



  // Create a unique job id for a table load.

  static String createJobId(

      String prefix, TableDestination tableDestination, int partition, long index) {

    // Job ID must be different for each partition of each table.

    String destinationHash =

        Hashing.murmur3_128().hashUnencodedChars(tableDestination.toString()).toString();

    String jobId = String.format("%s_%s", prefix, destinationHash);

    if (partition >= 0) {

      jobId += String.format("_%05d", partition);

    }

    if (index >= 0) {

      jobId += String.format("_%05d", index);

    }

    return jobId;

  }



  @VisibleForTesting

  static class JsonSchemaToTableSchema implements SerializableFunction<String, TableSchema> {

    @Override

    public TableSchema apply(String from) {

      return fromJsonString(from, TableSchema.class);

    }

  }



  static class TableSchemaToJsonSchema implements SerializableFunction<TableSchema, String> {

    @Override

    public String apply(TableSchema from) {

      return toJsonString(from);

    }

  }



  static class JsonTableRefToTableRef implements SerializableFunction<String, TableReference> {

    @Override

    public TableReference apply(String from) {

      return fromJsonString(from, TableReference.class);

    }

  }



  static class JsonTableRefToTableSpec implements SerializableFunction<String, String> {

    @Override

    public String apply(String from) {

      return toTableSpec(fromJsonString(from, TableReference.class));

    }

  }



  static class TableRefToJson implements SerializableFunction<TableReference, String> {

    @Override

    public String apply(TableReference from) {

      return toJsonString(from);

    }

  }



  static class TableRefToTableRefProto

      implements SerializableFunction<TableReference, TableReferenceProto.TableReference> {

    @Override

    public TableReferenceProto.TableReference apply(TableReference from) {

      return toTableRefProto(from);

    }

  }



  static class TableRefToTableSpec implements SerializableFunction<TableReference, String> {

    @Override

    public String apply(TableReference from) {

      return toTableSpec(from);

    }

  }



  static class TableRefProtoToTableSpec

      implements SerializableFunction<TableReferenceProto.TableReference, String> {

    @Override

    public String apply(TableReferenceProto.TableReference from) {

      return toTableSpec(from);

    }

  }



  @VisibleForTesting

  static class TableSpecToTableRef implements SerializableFunction<String, TableReference> {

    @Override

    public TableReference apply(String from) {

      return parseTableSpec(from);

    }

  }



  static class TimePartitioningToJson implements SerializableFunction<TimePartitioning, String> {

    @Override

    public String apply(TimePartitioning partitioning) {

      return toJsonString(partitioning);

    }

  }



  static String createJobIdToken(String jobName, String stepUuid) {

    return String.format("unified_job_%s_%s", stepUuid, jobName.replaceAll("-", ""));

  }



  static String getExtractJobId(String jobIdToken) {

    return String.format("%s-extract", jobIdToken);

  }



  static TableReference createTempTableReference(String projectId, String jobUuid) {

    String queryTempDatasetId = "temp_dataset_" + jobUuid;

    String queryTempTableId = "temp_table_" + jobUuid;

    return new TableReference()

        .setProjectId(projectId)

        .setDatasetId(queryTempDatasetId)

        .setTableId(queryTempTableId);

  }



  static String resolveTempLocation(

      String tempLocationDir, String bigQueryOperationName, String stepUuid) {

    return FileSystems.matchNewResource(tempLocationDir, true)

        .resolve(bigQueryOperationName, ResolveOptions.StandardResolveOptions.RESOLVE_DIRECTORY)

        .resolve(stepUuid, ResolveOptions.StandardResolveOptions.RESOLVE_DIRECTORY)

        .toString();

  }

}